import taichi as ti
from ..camera import Camera
from ..utils import *


class Fisheye(Camera):
    """鱼眼相机"""

    def __init__(
        self,
        size: "tuple[int, int]",
        fov: float = 60,
    ):
        super().__init__(size)
        self.fov: ti.ScalarField = ti.field(float, shape=1)
        self.fov[0] = fov
        self.xyz: ti.MatrixField = Mat3x3.field(shape=1)
        self.setScreen()

    def setScreen(self):
        sinTheta = ti.sin(self.theta)
        z = Vec3(
            ti.cos(self.phi) * sinTheta, ti.cos(self.theta), ti.sin(self.phi) * sinTheta
        )
        x = Vec3(ti.cos(self.phi - PI / 2), 0, ti.sin(self.phi - PI / 2))
        y = z.cross(x)
        self.x, self.y, self.z = Vec4(x, 0), Vec4(y, 0), Vec4(z, 0)
        self.xyz[0] = Mat3x3(x, y, z).transpose()
        return self

    @ti.func
    def rayAt(self, uv: Vec2, t: int = 0) -> Ray4:
        ss = superSampling(self.SSNumber, t)
        halfWH = Vec2(self.w, self.h) / 2
        shift = uv + ss - halfWH
        r = shift.norm() / halfWH.x
        theta, phi = r * self.fov[0] / 2 * deg2rad, ti.atan2(shift.y, shift.x)
        locXYZ = Vec3(
            ti.sin(theta) * ti.cos(phi), ti.sin(theta) * ti.sin(phi), ti.cos(theta)
        )
        return Ray4(self.origin[0], Vec4((self.xyz[0] @ locXYZ).normalized(), -1))

    def makeSettingPlane(self, gui: ti.ui.Gui) -> bool:
        clear = False
        if (new_value := gui.slider_float("fov", self.fov[0], 1, 360)) != self.fov[0]:
            self.fov[0] = new_value
            clear = True
        if clear:
            self.setScreen()
        return clear
